from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
from matplotlib.colors import TwoSlopeNorm

# analysis of phase space

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal

# 全局变量，用于存储和图像显示线程交互的数据
class imgview:
    global_image = None
    imgview_exit = False
    trajectory = []
    focus_i = 0
    traj_i = 0

imgview_data = imgview()

# 定义一个函数，用于在独立线程中显示图像
def show_image():
    grid_size_display = 20
    while not imgview_data.imgview_exit:
        # 检查全局变量是否有图像
        if imgview_data.global_image is not None:
            img = np.copy(imgview_data.global_image)
            state_x = imgview_data.trajectory[imgview_data.traj_i][0]
            state_y = imgview_data.trajectory[imgview_data.traj_i][1]
            cv2.circle(img, (state_y * grid_size_display + int(grid_size_display/2), state_x * grid_size_display + int(grid_size_display/2)), 7, (0, 0, 255), -1, cv2.LINE_AA)
            # 显示图像
            cv2.imshow("Image", img)
            key = cv2.waitKey(1)
            if key == ord('a'):
                imgview_data.focus_i -= 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
            elif key == ord('d'):
                imgview_data.focus_i += 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
        else:
            # 图像还未产生，等待100毫秒
            time.sleep(0.1)

def get_intrinsic_pc():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)
    
    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 1000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)
    rnn_state_old = rnn_state.copy()
    diff = jnp.abs(rnn_state - rnn_state_old)
    rnn_state_old = rnn_state.copy()
    diff_norm = jnp.linalg.norm(diff, axis=1)
    diff_norm_old = diff_norm.copy()
    norm_std = diff_norm.copy()

    rnn_state_init = rnn_state.copy()

    obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

    rnn_state_trajectory = []

    for t in range(rpl_config.life_duration):

        if t == rpl_config.probe_point:
            rnn_state_init = rnn_state.copy()

        progress_bar(t, rpl_config.life_duration)

        rnn_state_trajectory.append(np.array(rnn_state).copy())
        
        """ model forward 
        """
        rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
        
            
    print(rnn_state.shape)
    print(norm_std.shape)
    rnn_state_np = np.array(rnn_state)

    # 将 rnn_state_trajectory 展开成 rnn_state_np 的形状
    rnn_state_trajectory_np = np.array(rnn_state_trajectory)
    rnn_state_trajectory_np = rnn_state_trajectory_np.reshape(-1, rnn_state_trajectory_np.shape[-1])
    print("shape of rnn_state_trajectory_np: ", rnn_state_trajectory_np.shape)

    # 对 rnn_state_np 进行 PCA
    pca = PCA()
    # pca.fit(rnn_state_np)
    pca.fit(rnn_state_trajectory_np)

    # 打印 variance ratio
    print(pca.explained_variance_ratio_)

    rnn_state_np_pca = pca.transform(rnn_state_np)

    # 创建KMeans对象，指定聚类数为4
    kmeans = KMeans(n_clusters=4)
    # 对rnn_state_np_pca进行聚类
    kmeans.fit(rnn_state_np_pca)
    # 获取聚类中心的坐标
    cluster_centers = kmeans.cluster_centers_

    return pca, cluster_centers

def main():

    intrinsic_pca, cluster_centers = get_intrinsic_pc()
    # print(intrinsic_pca.explained_variance_ratio_)

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    landscapes, states, goals = [], [], []
    
    landscape_, state_, goal_ = load_task(rpl_config.task_pth, display=False)

    landscapes.append(landscape_)
    states.append(state_)
    goals.append(goal_)

    states = jnp.array(states)
    goals = jnp.array(goals)

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)

    bias1 = np.array(tree_leaves[0])
    mat1 = np.array(tree_leaves[1])
    print("mat1.shape: ", mat1.shape)

    mat_obs = np.array(tree_leaves[1])[128:137,:]
    mat_intr = np.array(tree_leaves[1])[0:128,:]
    print("mat_obs.shape: ", mat_obs.shape)

    
    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    """ create grid env
    """
    start_time = time.time()

    GE = GridEnv(landscapes = landscapes, width = 12, height = 12, num_envs_per_landscape = 1, reward_free=True)
    GE.reset()
    print("time taken to create envs: ", time.time() - start_time)

    # set states of GE
    GE.batched_states = states.copy()
    # set goals of GE
    GE.batched_goals = goals.copy()
    GE.init_batched_states, GE.init_batched_goals = jnp.copy(GE.batched_states), jnp.copy(GE.batched_goals)
    GE.batched_goal_reached = batch_compute_goal_reached(GE.batched_states, GE.batched_goals)
    GE.last_batched_goal_reached = jnp.copy(GE.batched_goal_reached)
    GE.concat_obs = get_ideal_obs_vmap(GE.batched_envs, GE.batched_states, GE.batched_goals, GE.last_batched_goal_reached)
    concat_obs = GE.concat_obs

    print("shape of concat_obs: ", concat_obs[0].shape)

    rnn_state = model.initial_state(GE.num_envs)
    step_count = 0
    trajectories = []
    goal_record = []
    HS_trajectory = []
    IPFs = []
    intr_field = []

    for t in range(rpl_config.life_duration):

        progress_bar(t, rpl_config.life_duration)

        step_count += 1

        # 将rnn_state[0]和concat_obs[0]拼接成一个新的向量
        new_vector = np.concatenate((rnn_state[0], concat_obs[0]))
        # 将新向量与mat1相乘，得到长度为128的向量
        result_vector = np.dot(new_vector, mat1) + bias1
        result_vector = np.tanh(result_vector)
        intr_vector = np.dot(rnn_state[0], mat_intr) + bias1
        intr_vector = np.tanh(intr_vector)

        IPF = result_vector - intr_vector
        IPFs.append(IPF)
        intr_field.append(intr_vector - rnn_state[0])

        trajectories.append(np.array(GE.batched_states[0]))
        HS_trajectory.append(np.array(rnn_state[0]))
        goal_record.append(GE.batched_goal_reached[0])

        """ model forward and step the env
        """
        rnn_state, y1 = model_forward(params, rnn_state, concat_obs, model)
        batched_actions = get_action_vmap(y1)
        batched_goal_reached, concat_obs = GE.step(batched_actions, reset=True)

    img = GE.render2()
    imgview_data.global_image = img
    imgview_data.trajectory = trajectories

    print("shape of trajectories: ", np.array(trajectories).shape)
    print("shape of goal_record: ", np.array(goal_record).shape)
    print("shape of rnn_state: ", rnn_state.shape)
    print("shape of HS_trajectory: ", np.array(HS_trajectory).shape)

    # # 重新计算pca
    # intrinsic_pca = PCA()
    # intrinsic_pca.fit(np.array(HS_trajectory))

    HS_pca = intrinsic_pca.transform(np.array(HS_trajectory))
    IPF_pca = intrinsic_pca.transform(np.array(IPFs) + intrinsic_pca.mean_)
    intr_pca = intrinsic_pca.transform(np.array(intr_field) + intrinsic_pca.mean_)

    # 将 pca 空间的点，逆变换回原始相空间
    HS_pca[:, 3:] = 0       # 将方差小的维度置零
    HS_trajectory_1 = intrinsic_pca.inverse_transform(HS_pca)
    # 比较 HS_1 和 HS_trajectory 的差异
    print("HS_1 - HS_trajectory: ", np.linalg.norm(HS_trajectory_1 - np.array(HS_trajectory)))

    # 计算主成分空间中的内禀向量场
    intr_field1 = []
    for i in range(HS_trajectory_1.shape[0]):
        intr_vector = np.dot(HS_trajectory_1[i], mat_intr) + bias1
        intr_vector = np.tanh(intr_vector)
        intr_field1.append(intr_vector - HS_trajectory_1[i])
    # 将内禀向量场投影到主成分空间
    intr_pca_1 = intrinsic_pca.transform(np.array(intr_field1) + intrinsic_pca.mean_)
    
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    # 计算每个点的颜色值
    colors = np.linspace(0, 1, len(HS_pca))
    # 绘制3D折线图，并为每个点设置颜色
    ax.scatter(cluster_centers[:, 0], cluster_centers[:, 1], cluster_centers[:, 2], c='b', s=140)
    # for i in range(len(HS_pca) - 1):
    #     ax.plot([HS_pca[i, 0], HS_pca[i+1, 0]], [HS_pca[i, 1], HS_pca[i+1, 1]], [HS_pca[i, 2], HS_pca[i+1, 2]], color=plt.cm.RdYlBu(colors[i]))

    for i in range(len(HS_pca) - 1):
        ax.quiver(HS_pca[i, 0], HS_pca[i, 1], HS_pca[i, 2], intr_pca[i, 0], intr_pca[i, 1], intr_pca[i, 2], color='g', length=0.2, arrow_length_ratio=0.3)
        ax.quiver(HS_pca[i, 0], HS_pca[i, 1], HS_pca[i, 2], intr_pca_1[i, 0], intr_pca_1[i, 1], intr_pca_1[i, 2], color='r', length=0.2, arrow_length_ratio=0.3)

    # for i in range(len(HS_pca) - 1):
    #     ax.quiver(HS_pca[i, 0], HS_pca[i, 1], 0, intr_pca[i, 0], intr_pca[i, 1], 0, color='g', length=0.2, arrow_length_ratio=0.3)
    #     ax.quiver(HS_pca[i, 0], HS_pca[i, 1], 0, intr_pca_1[i, 0], intr_pca_1[i, 1], 0, color='r', length=0.2, arrow_length_ratio=0.3)

    plt.show()

def ivf():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)

    bias1 = np.array(tree_leaves[0])

    mat_intr = np.array(tree_leaves[1])[0:128,:]
    print("mat_intr.shape: ", mat_intr.shape)

    # 生成一个包含随机个数的整数集合，元素个数大于20个
    num_indices = np.random.randint(21, 128)

    # 生成一个包含[0, 127]范围内所有整数的数组
    all_indices = np.arange(128)

    # 从中随机选择num_indices个不重复的元素
    random_indices = np.random.choice(all_indices, size=num_indices, replace=False)

    # random_indices = np.array([i for i in range(128)])

    # 从小到大排列
    random_indices.sort()

    # 输出结果
    print("num_indices: ", num_indices)
    print(random_indices)

    # 选择性保留mat_intr的第一个维度
    new_mat = np.take(mat_intr, random_indices, axis=0)
    new_mat = np.take(new_mat, random_indices, axis=1)
    new_bias = np.take(bias1, random_indices, axis=0)
    print("new_mat.shape: ", new_mat.shape)
    print("new_bias.shape: ", new_bias.shape)
    print("mat_intr.shape: ", mat_intr.shape)
    print("bias1.shape: ", bias1.shape)

    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 1000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)

    new_rnn_state = jnp.take(rnn_state, random_indices, axis=1)
    rnn_state = np.array(new_rnn_state)
    print("rnn_state.shape: ", rnn_state.shape)
    print("new_rnn_state.shape: ", new_rnn_state.shape)

    rnn_state_trajectory = []
    intr_field = []

    for t in range(rpl_config.life_duration):

        progress_bar(t, rpl_config.life_duration)

        intr_vector = np.dot(rnn_state, new_mat) + new_bias
        intr_vector = np.tanh(intr_vector)
        intr_field.append(intr_vector - rnn_state)

        rnn_state_trajectory.append(np.array(rnn_state).copy())
        
        rnn_state = intr_vector
        
    print(rnn_state.shape)
    rnn_state_np = np.array(rnn_state)

    rnn_state_trajectory_np = np.array(rnn_state_trajectory)

    # 将 rnn_state_trajectory_np 的第一个维度和第二个维度互换
    rnn_state_trajectory_np = rnn_state_trajectory_np.transpose(1, 0, 2)
    print("shape of rnn_state_trajectory_np: ", rnn_state_trajectory_np.shape)

    # 将rnn_state_trajectory_np沿着第二个维度分割成多个20步的窗口
    windows = np.array_split(rnn_state_trajectory_np, rnn_state_trajectory_np.shape[1] // 5, axis=1)

    # 计算每个窗口内每个独立序列的方差
    variances = np.var(windows, axis=2)

    # 计算每个窗口内每个独立序列的平均值
    means = np.mean(variances, axis=2)

    # 输出结果
    print("shape of means: ", means.shape)

    # 假设means是一个形状为(100, 100)的二维数组
    # 计算最大绝对值
    max_abs = np.max(np.abs(means))

    # 定义配色方案
    cmap = plt.cm.get_cmap('coolwarm')
    norm = TwoSlopeNorm(vmin=-max_abs, vcenter=0, vmax=max_abs)

    # 绘制二维色温图
    plt.imshow(means, cmap=cmap, norm=norm, interpolation='nearest')
    plt.colorbar()
    plt.show()




if __name__ == "__main__":
    
    # main()

    ivf()